from pycocotools.coco import COCO
import cv2 as cv
import os, sys


CLASSES = ('aeroplane', 'bicycle', 'bird', 'boat',
        'bottle', 'bus', 'car', 'cat', 'chair',
        'cow', 'diningtable', 'dog', 'horse',
        'motorbike', 'person', 'pottedplant',
        'sheep', 'sofa', 'train', 'tvmonitor')


def parse_json(annPath):
    '''
    res: List of ground truth
    annotations: dict of ground truth
    '''
    coco = COCO(annotation_file=annPath)
    imgIds = coco.getImgIds()
    annIds = coco.getAnnIds()
    catIds = coco.getCatIds()
    annotations = dict()

    for Id in annIds:
        annInfos = coco.loadAnns(Id)
        for info in annInfos:
            imageId = info['image_id']
            bbox = info['bbox']
            catId = info['category_id']
            imInfos = coco.loadImgs(imageId)
            if len(imInfos) > 1:
                print(2)

            for imginfo in imInfos:
                jpg = imginfo['file_name']
                if jpg not in annotations.keys():
                    annotations[jpg] = []
            annotations[jpg].append(bbox + [catId])
    res = []
    for k, v in annotations.items():
        res.append([k, v])
    return annotations, res


def visualize_bboxes(anns):
    for name, bboxes in anns.items():
        path = os.path.join('E:\dataset\VOC2012\VOCdevkit\VOC2012\JPEGImages', name)
        assert os.path.exists(path)
        im = cv.imread(path)
        for bbox in bboxes:
            cv.rectangle(im, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (244, 2, 2), 1)
            cv.putText(im, CLASSES[bbox[-1]], (bbox[0], bbox[1]), cv.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
        cv.imshow('d', im)
        cv.waitKey(0)
    cv.destroyAllWindows()


if __name__ == '__main__':
    annPath = './voc2012_train.json'
    anns, res = parse_json(annPath)
    visualize_bboxes(anns)














